#include "ompl/base/goals/GoalSampleableRegion.h"
#include "RRTIterationCount.h"

using namespace ompl;

RRTIterationCount::RRTIterationCount(const ompl::base::SpaceInformationPtr &si)
: ompl::geometric::RRT(si)
{
	// for DRAWING
	/*
	goalBias_ = 0.01;
	maxDistance_ = 5;
	*/
}

RRTIterationCount::~RRTIterationCount(void)
{
}

void RRTIterationCount::clear(void)
{
	ompl::geometric::RRT::clear();
	uiNumIterations_ = 0;
}

// for DRAWING
bool isCollisionFree(base::SpaceInformationPtr si_, base::State* from, base::State* to)
{
	
	std::vector<double> vState(3), vState2(3), v(3);
	si_->getStateSpace()->serialize(&(vState[0]), from);
	si_->getStateSpace()->serialize(&(vState2[0]), to);

	for (int i = 0; i <= 100; ++i)
	{
		v[0] = (i / 100.0) * vState[0] + (100 - i) / 100.0 * vState2[0];
		v[1] = (i / 100.0) * vState[1] + (100 - i) / 100.0 * vState2[1];
		if (v[0] > 20 || v[0] < -20)
			continue;
		if (v[1] > 20 || v[1] < -20)
			continue;

		return false;
	}
	
	return true;
}

// ++uiNumIterations_ is added
ompl::base::PlannerStatus RRTIterationCount::solve(const ompl::base::PlannerTerminationCondition &ptc)
{
	checkValidity();
    base::Goal                 *goal   = pdef_->getGoal().get();
    base::GoalSampleableRegion *goal_s = dynamic_cast<base::GoalSampleableRegion*>(goal);

    while (const base::State *st = pis_.nextStart())
    {
        Motion *motion = new Motion(si_);
        si_->copyState(motion->state, st);
        nn_->add(motion);
    }

    if (nn_->size() == 0)
    {
        OMPL_ERROR("There are no valid initial states!");
        return base::PlannerStatus::INVALID_START;
    }

    if (!sampler_)
        sampler_ = si_->allocStateSampler();

    OMPL_INFORM("Starting with %u states", nn_->size());

    Motion *solution  = NULL;
    Motion *approxsol = NULL;
    double  approxdif = std::numeric_limits<double>::infinity();
    Motion *rmotion   = new Motion(si_);
    base::State *rstate = rmotion->state;
    base::State *xstate = si_->allocState();

    while (ptc == false)
    {
		++uiNumIterations_;

        /* sample random state (with goal biasing) */
        if (goal_s && rng_.uniform01() < goalBias_ && goal_s->canSample())
            goal_s->sampleGoal(rstate);
        else
            sampler_->sampleUniform(rstate);

        /* find closest state in the tree */
        Motion *nmotion = nn_->nearest(rmotion);
        base::State *dstate = rstate;

        /* find state to add */
        double d = si_->distance(nmotion->state, rstate);
        if (d > maxDistance_)
        {
            si_->getStateSpace()->interpolate(nmotion->state, rstate, maxDistance_ / d, xstate);
            dstate = xstate;
        }

        if (si_->checkMotion(nmotion->state, dstate))
        {
            /* create a motion */
            Motion *motion = new Motion(si_);
            si_->copyState(motion->state, dstate);
            motion->parent = nmotion;

            nn_->add(motion);
            double dist = 0.0;
            bool sat = goal->isSatisfied(motion->state, &dist);
            if (sat)
            {
                approxdif = dist;
                solution = motion;
                break;
            }
            if (dist < approxdif)
            {
                approxdif = dist;
                approxsol = motion;
            }
        }
    }

    bool solved = false;
    bool approximate = false;
    if (solution == NULL)
    {
        solution = approxsol;
        approximate = true;
    }

    if (solution != NULL)
    {
        lastGoalMotion_ = solution;

        /* construct the solution path */
        std::vector<Motion*> mpath;
        while (solution != NULL)
        {
            mpath.push_back(solution);
            solution = solution->parent;
        }

        /* set the solution path */
        ompl::geometric::PathGeometric *path = new ompl::geometric::PathGeometric(si_);
           for (int i = mpath.size() - 1 ; i >= 0 ; --i)
            path->append(mpath[i]->state);
        pdef_->addSolutionPath(base::PathPtr(path), approximate, approxdif);
        solved = true;
    }

    si_->freeState(xstate);
    if (rmotion->state)
        si_->freeState(rmotion->state);
    delete rmotion;

    OMPL_INFORM("Created %u states", nn_->size());

	// for DRAWING
	/*
	std::vector<Motion*> tree;
	nn_->list(tree);
	std::vector<double> vState1, vState2;
	vState1.resize(3);
	vState2.resize(3);
	for (int i = 0; i < tree.size(); ++i)
	{
		if (tree[i]->parent)
		{
			si_->getStateSpace()->serialize(&(vState1[0]), tree[i]->state);
			si_->getStateSpace()->serialize(&(vState2[0]), tree[i]->parent->state);
			std::cout << vState1[0] << " " << vState1[1] << " " << vState2[0] << " " << vState2[1] << std::endl;
		}
	}
	*/

    return base::PlannerStatus(solved, approximate);
}